from langchain_core.documents import Document
from langchain_milvus import Milvus
from pymilvus import connections, db, MilvusClient, DataType, CollectionSchema, FieldSchema
from core.FixedMilvus import FinalFixedMilvus
from core.config import settings
from core.OllamaEmbeddings_Create import get_embeddings  # 导入你的嵌入模型


def singleton(cls):
    _instance = {}

    def wrapper(*args, **kwargs):
        if cls not in _instance:
            _instance[cls] = cls(*args, **kwargs)
        return _instance[cls]

    return wrapper


@singleton
class GET_Milvus:
    def __init__(self):
        self._connect()  # 建立连接
        self.client = MilvusClient(uri=settings.MILVUS__HOST_AND_PORT)
        self.vector_store = None  # 向量存储实例
        self.default_dim = 768  # 默认向量维度
        self.FIXED_DB_NAME = "knowledge_db"  # 固定数据库名
        self.FIXED_COLLECTION_NAME = "knowledge_collection"  # 固定集合名
        # 向量维度：优先用配置，无配置则默认1024
        self.FIXED_VECTOR_DIM = getattr(settings, "VECTOR_DIMENSION", 1024)
        # 明确字段名（与集合Schema严格一致）
        self.TEXT_FIELD_NAME = "page_content"  # 文本字段名
        self.VECTOR_FIELD_NAME = "embedding"  # 向量字段名
        self.PRIMARY_FIELD_NAME = "id"  # 主键字段名
        self.embeddings = get_embeddings()
        sample_embedding = self.embeddings.embed_query("测试维度")
        self.FIXED_VECTOR_DIM = len(sample_embedding)  # 动态获取实际维度
        print(f"嵌入模型实际维度: {self.FIXED_VECTOR_DIM}")

    def _connect(self):
        """建立与Milvus的连接（全局唯一）"""
        try:
            if not connections.has_connection("default"):
                connections.connect(
                    alias="default",
                    host=settings.MILVUS__HOST,
                    port=settings.MILVUS__PORT
                )
            print("Milvus连接已建立（全局唯一）")
            return True
        except Exception as e:
            print(f"Milvus连接失败: {str(e)}")
            return False

    def init_fixed_collection(self, force_reinit: bool = False):
        """初始化固定集合（核心方法，确保字段配置正确）"""
        try:
            # 1. 切换到目标数据库
            self.use_database(self.FIXED_DB_NAME)  # 先尝试切换，失败则创建
            if not self._check_database_exists(self.FIXED_DB_NAME):
                self.create_database(self.FIXED_DB_NAME)
                self.use_database(self.FIXED_DB_NAME)

            # 2. 强制重建集合（确保字段配置生效）
            if force_reinit or not self.client.has_collection(self.FIXED_COLLECTION_NAME):
                if force_reinit:
                    self.delete_collection(self.FIXED_COLLECTION_NAME)
                # 关键：创建包含page_content的严格Schema（关闭动态字段）
                self._create_fixed_schema_collection()
                print(f"集合 {self.FIXED_COLLECTION_NAME} 创建/重建成功，字段配置已更新")

            # 3. 创建向量索引（提升搜索性能）
            self._create_vector_index()

            # 4. 初始化LangChain向量存储（绑定正确字段）
            self.vector_store = FinalFixedMilvus(
                embedding_function=get_embeddings(),
                connection_args={"uri": settings.MILVUS__HOST_AND_PORT},
                collection_name=self.FIXED_COLLECTION_NAME,
                primary_field=self.PRIMARY_FIELD_NAME,
                vector_field=self.VECTOR_FIELD_NAME,
                text_field=self.TEXT_FIELD_NAME,
                # 关键：通过构造函数指定需要返回的字段（父类会自动处理，不重复）
                # output_fields=[self.PRIMARY_FIELD_NAME, self.TEXT_FIELD_NAME]
            )

            # 验证字段配置
            self._verify_collection_fields()
            print(f"初始化完成，text_field: {self.vector_store._text_field}")
            return self.vector_store
        except Exception as e:
            print(f"初始化失败: {str(e)}")
            return None

    def _create_fixed_schema_collection(self):
        """创建固定Schema的集合（关闭动态字段，明确page_content属性）"""
        # 定义字段（严格指定类型和属性）
        fields = [
            # 主键字段（自增）
            FieldSchema(
                name=self.PRIMARY_FIELD_NAME,
                dtype=DataType.INT64,
                is_primary=True,
                auto_id=True  # 自动生成主键
            ),
            # 向量字段
            FieldSchema(
                name=self.VECTOR_FIELD_NAME,
                dtype=DataType.FLOAT_VECTOR,
                dim=self.FIXED_VECTOR_DIM  # 与嵌入模型维度一致
            ),
            # 文本字段（关键：明确VARCHAR类型和最大长度）
            FieldSchema(
                name=self.TEXT_FIELD_NAME,
                dtype=DataType.VARCHAR,
                max_length=65535  # 足够存储长文本
            )
        ]

        # 创建Schema（关闭动态字段，避免字段混乱）
        schema = CollectionSchema(
            fields=fields,
            description="固定Schema集合，包含page_content文本字段",
            enable_dynamic_field=False  # 关键：关闭动态字段，只允许定义的字段
        )

        # 创建集合
        self.client.create_collection(
            collection_name=self.FIXED_COLLECTION_NAME,
            schema=schema
        )

    def _create_vector_index(self):
        """为向量字段创建索引（确保搜索可用）"""
        try:
            # 检查索引是否已存在
            indexes = self.client.describe_index(
                collection_name=self.FIXED_COLLECTION_NAME,
                field_name=self.VECTOR_FIELD_NAME
            )
            if not indexes:
                # 创建IVF_FLAT索引（适合余弦相似度）
                index_params = {
                    "index_type": "IVF_FLAT",
                    "metric_type": "COSINE",
                    "params": {"nlist": 128}
                }
                self.client.create_index(
                    collection_name=self.FIXED_COLLECTION_NAME,
                    field_name=self.VECTOR_FIELD_NAME,
                    index_params=index_params
                )
                print(f"向量字段 {self.VECTOR_FIELD_NAME} 索引创建成功")
            else:
                print(f"向量字段 {self.VECTOR_FIELD_NAME} 索引已存在")
        except Exception as e:
            print(f"创建向量索引失败: {str(e)}")

    def _verify_collection_fields(self):
        """验证集合字段是否包含page_content"""
        try:
            collection_info = self.client.describe_collection(self.FIXED_COLLECTION_NAME)
            fields = [f["name"] for f in collection_info["schema"]["fields"]]
            if self.TEXT_FIELD_NAME not in fields:
                raise ValueError(f"集合 {self.FIXED_COLLECTION_NAME} 中缺少 {self.TEXT_FIELD_NAME} 字段")
            print(f"集合字段验证通过，包含: {fields}")
        except Exception as e:
            print(f"集合字段验证失败: {str(e)}")

    def _check_database_exists(self, db_name: str):
        """检查数据库是否存在"""
        try:
            dbs = db.list_database()
            return db_name in dbs
        except Exception as e:
            print(f"检查数据库存在性失败: {str(e)}")
            return False

    # 数据库操作
    def create_database(self, db_name: str):
        try:
            db.create_database(db_name)
            print(f"数据库 {db_name} 创建成功")
            return True
        except Exception as e:
            print(f"创建数据库失败: {str(e)}")
            return False

    def use_database(self, db_name: str):
        try:
            db.using_database(db_name)
            print(f"已切换到数据库 {db_name}")
            return True
        except Exception as e:
            print(f"切换数据库失败: {str(e)}")
            return False

    def delete_database(self, db_name: str):
        try:
            db.drop_database(db_name)
            print(f"数据库 {db_name} 删除成功")
            return True
        except Exception as e:
            print(f"删除数据库失败: {str(e)}")
            return False

    # 集合操作
    def delete_collection(self, collection_name: str):
        try:
            if self.client.has_collection(collection_name):
                self.client.drop_collection(collection_name)
                print(f"集合 {collection_name} 删除成功")
                return True
            print(f"集合 {collection_name} 不存在")
            return False
        except Exception as e:
            print(f"删除集合失败: {str(e)}")
            return False

    def list_collections(self):
        try:
            return self.client.list_collections()
        except Exception as e:
            print(f"列出集合失败: {str(e)}")
            return []

    # 数据插入（确保page_content字段正确）
    def insert_documents(self, documents: list):
        """插入文档数据（确保page_content正确写入）"""
        try:
            milvus_data = []
            for doc in documents:
                # 生成向量（如果文档中没有预计算的embedding）
                if isinstance(doc, Document):
                    embedding = self.embeddings.embed_query(doc.page_content)  # 用嵌入模型生成向量
                    milvus_data.append({
                        self.VECTOR_FIELD_NAME: embedding,
                        self.TEXT_FIELD_NAME: doc.page_content  # 明确写入page_content
                    })
                else:
                    # 处理字典格式
                    embedding = doc.get(self.VECTOR_FIELD_NAME) or self.embeddings.embed_query(
                        doc.get(self.TEXT_FIELD_NAME, ""))
                    milvus_data.append({
                        self.VECTOR_FIELD_NAME: embedding,
                        self.TEXT_FIELD_NAME: doc.get(self.TEXT_FIELD_NAME, "")  # 确保有默认值
                    })

            res = self.client.insert(
                collection_name=self.FIXED_COLLECTION_NAME,
                data=milvus_data
            )
            print(f"插入成功，数量: {res['insert_count']}")
            return res
        except Exception as e:
            print(f"插入文档失败: {str(e)}")
            return None

    # 搜索验证（检查page_content是否返回）
    def verify_search_page_content(self, query_vector: list):
        """用原生客户端验证搜索是否返回page_content"""
        try:
            res = self.client.search(
                collection_name=self.FIXED_COLLECTION_NAME,
                data=[query_vector],
                limit=1,
                output_fields=[self.PRIMARY_FIELD_NAME, self.TEXT_FIELD_NAME]  # 只返回需要的字段
            )
            if res and len(res) > 0 and len(res[0]) > 0:
                entity = res[0][0].entity.to_dict()
                print(f"原生搜索返回的字段: {list(entity.keys())}")
                if self.TEXT_FIELD_NAME in entity:
                    print(f"page_content内容: {entity[self.TEXT_FIELD_NAME][:50]}...")  # 打印前50字符
                    return True
                else:
                    print(f"原生搜索未返回 {self.TEXT_FIELD_NAME}")
                    return False
            else:
                print("搜索无结果")
                return False
        except Exception as e:
            print(f"验证搜索失败: {str(e)}")
            return False

    def __del__(self):
        """关闭连接"""
        if connections.has_connection("default"):
            connections.disconnect("default")
            print("Milvus连接已关闭")